'''
Description: 可视化点云
Author: suyunzheng
Date: 2021-10-20 22:40:40
LastEditTime: 2021-10-21 15:40:21
LastEditors: maple
'''
import sys
import os
import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt

import yaml
import glob
import logging


logging.basicConfig(level=logging.debug, filename="demo.log", filemode="w", 
                    format ="%(asctime)s - %(name)s - %(levelname)-9s - %(filename)-8s : %(lineno)s line - %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S"  )


COLOR_MAP = {
    0: (0, 0., 0.),
    1: (174., 199., 232.),
    2: (152., 223., 138.),
    3: (31., 119., 180.),
    4: (255., 187., 120.),
    5: (188., 189., 34.),
    6: (140., 86., 75.),
    7: (255., 152., 150.),
    8: (214., 39., 40.),
    9: (197., 176., 213.),
    10: (148., 103., 189.),
    11: (196., 156., 148.),
    12: (23., 190., 207.),
    13: (100., 85., 144.),
    14: (247., 182., 210.),
    15: (66., 188., 102.),
    16: (219., 219., 141.),
    17: (140., 57., 197.),
    18: (202., 185., 52.),
    19: (51., 176., 203.),
    20: (200., 54., 131.),
    21: (92., 193., 61.),
    22: (78., 71., 183.),
    23: (172., 114., 82.),
    24: (255., 127., 14.),
    25: (91., 163., 138.),
    26: (153., 98., 156.),
    27: (140., 153., 101.),
    28: (158., 218., 229.),
    29: (100., 125., 154.),
    30: (178., 127., 135.),
    32: (146., 111., 194.),
    33: (44., 160., 44.),
    34: (112., 128., 144.),
    35: (96., 207., 209.),
    36: (227., 119., 194.),
    37: (213., 92., 176.),
    38: (94., 106., 211.),
    39: (82., 84., 163.),
    40: (100., 85., 144.),
}


def visualize_with_label(cloud, label, window_name="open3d"):
    # if label.shape == 1:
    # label = label.reshape((-1,1))
    assert cloud.shape[0] == label.shape[0]

    pt = o3d.geometry.PointCloud()
    pt.points = o3d.utility.Vector3dVector(cloud)
    colors = [COLOR_MAP[i] for i in list(label[:,0])]
    colors = np.asarray(colors).reshape((-1,3))/255
    pt.colors = o3d.utility.Vector3dVector(colors)

    o3d.visualization.draw_geometries([pt], window_name ,width=500, height=500)


def visualize_without_label(cloud, window_name="open3d"):
    # assert cloud.shape[0] == label.shape[0]

    pt = o3d.geometry.PointCloud()
    pt.points = o3d.utility.Vector3dVector(cloud)
    # colors = [COLOR_MAP[i] for i in list(label)]
    # pt.colors = o3d.utility.Vector3dVector(colors)

    o3d.visualization.draw_geometries([pt], window_name ,width=500, height=500)

'''
description: 可视化点云矩阵
param  {*}
return {*}
param {*} data: n*3
param {*} label: n*1
param {*} window_name
'''
def vis(data,label, window_name ):
    '''
    :param data: n*3的矩阵
    :param label: n*1的矩阵
    :return: 可视化
    '''
    data=data[:,:3]
    labels=np.asarray(label)
    max_label=labels.max()

    # 颜色
    # colors = plt.get_cmap("tab20")(labels / (max_label if max_label > 0 else 1))
    colors = plt.get_cmap("tab20")(labels / 19 )

    pt1 = o3d.geometry.PointCloud()
    pt1.points = o3d.utility.Vector3dVector(data.reshape(-1, 3))
    pt1.colors=o3d.utility.Vector3dVector(colors[:, 0, :3])

    o3d.visualization.draw_geometries([pt1], window_name ,width=500,height=500)


'''
description: 可视化点云数据流模块
param  {*}
return {*}
'''
class vis_pt_stream:
    def __init__(self, pt_dir) -> None:
        self.pt_dir = pt_dir
        self.filelist = []

        with open("/home/ubuntu/suyunzheng_ws/code/PolarSeg/semantic-kitti.yaml", 'r') as stream:
            semkittiyaml = yaml.safe_load(stream)
        self.learning_map = semkittiyaml['learning_map']            # 类别映射
        self.learning_map_inv = semkittiyaml['learning_map_inv']    # 颜色反映射    
        self.color_map = semkittiyaml['color_map']                  # 颜色设置
        self.return_ref = False

    def set_filelist(self) -> None:
        bin_name = os.path.join(self.pt_dir, "*.bin")
        for file in glob.glob(bin_name):
            self.filelist.append(file)
            self.filelist.sort()
            # logging.warning("file: {}".format(file))
        logging.warning("filelist: {}".format(len(self.filelist)))

    def dynamic_vis(self, hasLabel=False) -> None:
        vis = o3d.visualization.Visualizer()
        vis.create_window()
        pointcloud = o3d.geometry.PointCloud()
        to_reset = True
        vis.add_geometry(pointcloud)
        for f in self.filelist:
            logging.warning("file name: {}".format(f))
            pcd_arr,label_arr = self.get_pt_from_bin(f, hasLabel = hasLabel)
            pcd_arr.reshape((-1,3))
            label_arr.reshape((-1,1))
            # inverse_map 得到原始颜色
            raw_label_arr = np.vectorize(self.learning_map_inv.__getitem__)(label_arr)        # 类别映射
            color_map = [self.color_map[i] for i in list(raw_label_arr[:,0])]                 # 颜色映射
            colors = np.asarray(color_map).reshape((-1,3))

            pointcloud.points = o3d.utility.Vector3dVector(pcd_arr)
            pointcloud.colors = o3d.utility.Vector3dVector(colors)
            vis.update_geometry(pointcloud)
            if to_reset:
                vis.reset_view_point(True)
                to_reset = False
            vis.poll_events()
            vis.update_renderer()

    def get_pt_from_bin(self, filename, hasLabel=False, return_ref = False):
        if return_ref:
            self.return_ref = True
        raw_data = np.fromfile(filename, dtype=np.float32).reshape((-1, 4))
        if hasLabel:
            annotated_data = np.fromfile(filename.replace('velodyne','labels')[:-3]+'label', dtype=np.int32).reshape((-1,1))
            annotated_data = annotated_data & 0xFFFF #delete high 16 digits binary
            annotated_data = np.vectorize(self.learning_map.__getitem__)(annotated_data)        # 类别映射
        else:
            annotated_data = np.expand_dims(np.zeros_like(raw_data[:,0],dtype=int),axis=1)

        data_tuple = (raw_data[:,:3], annotated_data.astype(np.uint8))          # xyz
        if self.return_ref:
            data_tuple += (raw_data[:,3],)
        return data_tuple
        pass
        


if __name__=='__main__':
    ob1 = vis_pt_stream("/home/ubuntu/suyunzheng_ws/code/PolarSeg/data/sequences/08/velodyne")
    ob1.set_filelist()
    ob1.dynamic_vis(hasLabel=True)

